import tensorflow
from tensorflow.python.keras.callbacks import EarlyStopping, TensorBoard

import kashgari
from kashgari.tasks.labeling import BiGRU_CRF_Model
from kashgari.embeddings import BERTEmbedding


def read_file(path):
    features = [[]]
    labels = [[]]
    with open(path, 'r', encoding='utf-8') as f:
        message = f.readlines()
    for i in message:
        if i is not None:
            i = i.split("\t")
            if len(i) == 2:
                features.append(i[0].split())
                labels.append(i[1].split())
    return features, labels


train_data, train_labels = read_file("medicinal_data_indication/train_data.txt")
test_data, test_labels = read_file("medicinal_data_indication/test_data.txt")
valid_data, valid_labels = read_file("medicinal_data_indication/valid_data.txt")

# bert_embed = BERTEmbedding('wwm',
#                            trainable=True,
#                            task=kashgari.LABELING,
#                            sequence_length=10)
model = BiGRU_CRF_Model()
early_stopping = EarlyStopping(
    monitor='val_loss',
    min_delta=0,
    patience=8,
    verbose=1,
    mode='auto'
)
tb_cb = TensorBoard(log_dir="log_dir", write_images=1, histogram_freq=1)

model.fit(train_data,
          train_labels,
          valid_data,
          valid_labels,
          epochs=500,
          batch_size=256,
          callbacks=[early_stopping, tb_cb]
          )
model.evaluate(test_data, test_labels)
model.save("Indication")
